import os
import argparse

import cv2
import numpy as np
from tqdm import tqdm
import sys
import glfw

from cyw_devkit import tranform_pt_pad
import open3d as o3d
o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Debug)  # 开启调试信息

sys.path.append('../')  # 确保这个路径是正确的
from cyw_devkit.core.dataset import CywDataset, transform_bboxes
from cyw_devkit.core.visualizer import view3d
from cyw_devkit.core.visualizer import just_look_array, draw_pts2img, Visualizer
from tictoc import TicToc
import multiprocessing
from cyw_devkit.core.visualizer.open3d_arrow import get_yaw_arrows, get_speed_arrows, get_move_arrows
from matplotlib import pyplot as plt

# 配置参数
BEV_WIDTH = 800  # BEV图像宽度
BEV_HEIGHT = 800  # BEV图像高度
BEV_RESOLUTION = 0.2  # 每像素代表的实际距离（米/像素）
MAX_DEPTH = 50  # 最大显示深度（米）


def create_bev_map(points, boxes, velocities=None):
    # 创建空白BEV图像
    bev_image = np.zeros((BEV_HEIGHT, BEV_WIDTH, 3), dtype=np.uint8)

    # ========== 点云投影 ==========
    if len(points) > 0:
        # 转换到BEV坐标系（X向右，Y向下）
        u = (points[:, 0] / BEV_RESOLUTION + BEV_WIDTH // 2).astype(int)
        v = (points[:, 1] / BEV_RESOLUTION + BEV_HEIGHT // 2).astype(int)

        # 过滤超出边界的点
        mask = (u >= 0) & (u < BEV_WIDTH) & (v >= 0) & (v < BEV_HEIGHT)
        u, v = u[mask], v[mask]

        # 根据高度着色（蓝->红）
        heights = points[mask, 2]
        colors = plt.cm.jet((heights + 5) / 10)[:, :3] * 255  # 假设高度范围-5~5米

        # 绘制点云
        bev_image[v, u] = colors

    # ========== 绘制3D框 ==========
    for box in boxes:
        # 解析框参数（假设格式：[x, y, z, l, w, h, yaw, ...]）
        x, y, z, l, w, h, yaw = box[:7]

        # 计算框的四个角点（BEV视角）
        corners = np.array([
            [x + l / 2, y + w / 2],
            [x + l / 2, y - w / 2],
            [x - l / 2, y - w / 2],
            [x - l / 2, y + w / 2]
        ])

        # 应用旋转
        rot_mat = np.array([
            [np.cos(yaw), -np.sin(yaw)],
            [np.sin(yaw), np.cos(yaw)]
        ])
        corners = np.dot(corners - [x, y], rot_mat) + [x, y]

        # 转换到图像坐标
        img_corners = (corners[:, ::-1] / BEV_RESOLUTION + np.array([BEV_WIDTH // 2, BEV_HEIGHT // 2])).astype(int)

        # 绘制旋转矩形
        cv2.polylines(bev_image, [img_corners], True, (0, 255, 0), 2)

        # # 绘制方向箭头（车头方向）
        # arrow_len = max(l, w) * 0.8
        # end_point = np.array([
        #     x + np.cos(yaw) * arrow_len,
        #     y + np.sin(yaw) * arrow_len
        # ])
        # img_end = (end_point[::-1] / BEV_RESOLUTION + np.array([BEV_WIDTH // 2, BEV_HEIGHT // 2])).astype(int)
        # cv2.arrowedLine(bev_image,
        #                 tuple(img_corners[0].tolist()),
        #                 tuple(img_end.tolist()),
        #                 (0, 0, 255), 2)

    # ========== 绘制速度箭头 ==========
    if velocities is not None:
        for box, vel in zip(boxes, velocities):
            x, y = box[0], box[1]
            vx, vy = vel[0], vel[1]

            # 计算箭头终点
            speed = np.hypot(vx, vy)
            if speed > 0.1:  # 忽略静止目标
                scale = 10.0  # 速度缩放系数
                end_x = x + vx * scale
                end_y = y + vy * scale

                # 转换到图像坐标
                start = (int(x / BEV_RESOLUTION + BEV_WIDTH // 2),
                         int(y / BEV_RESOLUTION + BEV_HEIGHT // 2))
                end = (int(end_x / BEV_RESOLUTION + BEV_WIDTH // 2),
                       int(end_y / BEV_RESOLUTION + BEV_HEIGHT // 2))

                # 绘制速度箭头
                cv2.arrowedLine(bev_image, start, end, (255, 255, 0), 2,
                                tipLength=0.3)

    return bev_image

def main(args):
    # 统计耗时
    cost = TicToc("数据可视化")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    # for idx in range(process_size):
    #     main_worker(files[0], args)

    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(dataset_path, args):
    # cyw = CywDataset(dataset_path=dataset_path)
    cyw = CywDataset(dataset_path=dataset_path, radar_dim=args.radar_dim, motion=True)
    # output_file = os.path.join(args.out_path,os.path.basename(dataset_path)+'.mp4')
    # fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 设置视频编码格式（这里使用 MP4V 编码）
    # out = cv2.VideoWriter(output_file, fourcc, 30, (1080, 720))

    init_stamp = 0 # us
    boxes_all = []
    pc_all = []
    # vis = Visualizer(frame_data['lidar'].data, center_mode=False)
    # vis = Visualizer(points=None, center_mode=False)
    for idx, frame_data in enumerate(tqdm(cyw)):
        if idx<30:
            continue
            init_stamp = frame_data['lidar'].header.stamp
        #     vis = Visualizer(frame_data['lidar'].data, center_mode=False)

        # if idx<20 or idx%5!=0:
        #     continue
        # if idx>80:
        #     break

        if idx in cyw.keyframe_idxs: #判断关键帧

            boxes_now = frame_data['label'].data

            # # 生成BEV图像
            # bev_image = create_bev_map(
            #     points=frame_data['lidar'].data[:, :3],  # 使用x,y,z坐标
            #     boxes=boxes_now.tensors,
            #     velocities=boxes_now.tensors[:, 7:9] if boxes_now.tensors.shape[1] >= 9 else None
            # )
            #
            # # 显示结果
            # cv2.imshow("BEV View", bev_image)
            # cv2.waitKey(0)

            glfw.init()

            vis = Visualizer(frame_data['lidar'].data, center_mode=False)
            vis.add_bboxes(bbox3d=boxes_now.tensors)

            # if show_yaw and boxes_now.shape[1] >= 7:
            #     arrows = get_yaw_arrows(pre_box_set=boxes_now)
            #     for arrow in arrows:
            #         vis.o3d_visualizer.add_geometry(arrow)

            if boxes_now.tensors.shape[1] >= 9:
                arrows = get_speed_arrows(pre_box_set=boxes_now.tensors)
                for arrow in arrows:
                    vis.o3d_visualizer.add_geometry(arrow)
            vis.show()



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/home/adt/datasets/motion",
                        help='your data root')
    parser.add_argument('--radar-dim', type=int, default=5,
                        help='cyw 5;carla 7')
    parser.add_argument('--out-path', default="/home/adt/datasets/motion/dataview",
                        help='')
    parser.add_argument('--sync-time', action='store_true',
                        help='')
    args = parser.parse_args()

    main(args)
